import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
import xml.etree.ElementTree as ET
from glob import glob
from scipy.ndimage import gaussian_filter
import json


def json_points(label):
    with open(label, 'r') as f:
        info = json.load(f)
    points = info['points']
    return points
    
def get_result(pt_file, shape):
    density = torch.load(pt_file)
    count = "{:.2f}".format(density.sum().item())

    resize_density = F.interpolate(density, size=shape, mode="bilinear")
    
    x_sum = torch.sum(density, dim=(-1, -2))
    scale_factor = torch.nan_to_num(torch.sum(resize_density, dim=(-1, -2)) / x_sum, nan=0.0, posinf=0.0, neginf=0.0)
    resize_density = resize_density * scale_factor
    
    return count, resize_density.cpu().squeeze()

def generate_density_map(shape, label, sigma = 8):
    """
    Generate the density map based on the dot annotations provided by the label.
    """
    height, width = shape
    label = torch.tensor(label).float()
    
    density_map = torch.zeros((1, height, width), dtype=torch.float32)

    if len(label) > 0:
        assert len(label.shape) == 2 and label.shape[1] == 2, f"label should be a Nx2 tensor, got {label.shape}."
        label_ = label.long()
        label_[:, 0] = label_[:, 0].clamp(min=0, max=width - 1)
        label_[:, 1] = label_[:, 1].clamp(min=0, max=height - 1)
        density_map[0, label_[:, 1], label_[:, 0]] = 1.0

    if sigma is not None:
        assert sigma > 0, f"sigma should be positive if not None, got {sigma}."
        density_map = torch.from_numpy(gaussian_filter(density_map, sigma=sigma))

    return density_map.squeeze()

def LMDS_counting(input):
    input_max = torch.max(input).item()

    keep = nn.functional.max_pool2d(input, (3, 3), stride=1, padding=1)
    keep = (keep == input).float()
    input = keep * input

    '''set the pixel valur of local maxima as 1 for counting'''
    input[input < 100.0 / 255.0 * input_max] = 0
    input[input > 0] = 1

    ''' negative sample'''
    if input_max < 0.1:
        input = input * 0

    count = int(torch.sum(input).item())

    return count






def save_rgbtcc_rgb(names):
    savepath = "./05_save_rgbtcc_rgb"
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    
    gt_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'
    image_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'
    
    results = {
        'CLIP-EBC': '/data/store1/nzd/tir_cc/methods/clip/1_adapt/rgbtcc_results',
        'BL': '/data/store1/nzd/tir_cc/methods/bayesian/1_adapt_rgbtcc/results',
        'P2P': '/data/store1/nzd/tir_cc/methods/p2p/p2p_rgbtcc/results',
        'HMoDE': '/data/store1/nzd/tir_cc/methods/hmode/RMS_rgbtcc/results',
        'STEERER': '/data/store1/nzd/tir_cc/methods/steerer/STEERER_rgbtcc/results',
        'FIDTM': '/data/store1/nzd/tir_cc/methods/fidtm/25_optim_rgbtcc/results',
    }

    # 添加文字参数设置
    text_params = {
        's': "",      # 显示的文字内容
        'color': 'white',           # 文字颜色
        'fontfamily': 'serif',  # 字体
        'fontsize': 24,             # 字体大小（可调节）
        'ha': 'center',             # 水平居中
        'va': 'bottom',             # 垂直对齐底部
        'transform': None,  # 使用相对坐标系
    }

    for name in names:
        name = f"{name}R"
        img_file = f"{image_root}/{name}.jpg".replace('R.jpg', '_T.jpg')
        img = plt.imread(img_file)
        shape = img.shape[:2]
        
        gt_file = f"{gt_root}/{name}.xml".replace('R.xml', '_GT.json')
        points = json_points(gt_file)
        density = generate_density_map(shape, points)
    
        fig, ax = plt.subplots()
        ax.imshow(img)
        ax.imshow(density, cmap="jet", alpha=0.7)
        ax.axis('off')

        text_params['s'] = f"GT: {len(points)}"
        text_params['transform'] = ax.transAxes
        ax.text(x=0.5, y=0.02, **text_params) # 距离底部2%的位置（可调节） # 水平居中位置（50%）

        # plt.show()
        plt.savefig(f"{savepath}/{name}_density_{len(points)}.jpg", bbox_inches='tight', pad_inches=0, dpi=300)
        plt.close()
    
        for method, path in results.items():
            if method == 'P2P':
                pre_points = torch.load(os.path.join(path, f'{name}.pt')).detach().cpu().numpy().tolist()
                pre_count = "{:.2f}".format(len(pre_points))
                pre_density = generate_density_map(shape, pre_points, sigma=8)
            elif method == 'FIDTM':
                _, pre_density = get_result(os.path.join(path, f'{name}.pt'), shape)
                pre_count = LMDS_counting(torch.load(os.path.join(path, f'{name}.pt')))
                pre_count = "{:.2f}".format(pre_count)
            else:
                pre_count, pre_density = get_result(os.path.join(path, f'{name}.pt'), shape)
    
            if method == 'HMoDE':
                pre_count = "{:.2f}".format(float(pre_count) / 100.)

            fig, ax = plt.subplots()
            ax.imshow(img)
            ax.imshow(pre_density, cmap="jet", alpha=0.7)
            ax.axis('off')
            
            text_params['s'] = f"ET: {pre_count}"
            text_params['transform'] = ax.transAxes
            ax.text(x=0.5, y=0.02, **text_params) # 距离底部2%的位置（可调节） # 水平居中位置（50%）
            
            plt.savefig(f"{savepath}/{name}_{method}_{pre_count}.jpg", bbox_inches='tight', pad_inches=0, dpi=300)
            plt.close()






def save_rgbtcc_rgbt(names):
    savepath = "./05_save_rgbtcc_rgbt"
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    
    gt_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'
    image_root = '/data/store1/nzd/tir_cc/datasets/rgbtcc/test'
    
    results = {
        'MIANet': './MIANet_rgbtcc_results',
        'DLFIA': './DLFIA_rgbtcc_results',
        'FIDTM': '/data/store1/nzd/tir_cc/methods/fidtm/25_optim_rgbtcc/results',
    }


    # 添加文字参数设置
    text_params = {
        's': "",      # 显示的文字内容
        'color': 'white',           # 文字颜色
        'fontfamily': 'serif',  # 字体
        'fontsize': 24,             # 字体大小（可调节）
        'ha': 'center',             # 水平居中
        'va': 'bottom',             # 垂直对齐底部
        'transform': None,  # 使用相对坐标系
    }

    for name in names:
        name = f"{name}R"
        img_file = f"{image_root}/{name}.jpg".replace('R.jpg', '_T.jpg')
        img = plt.imread(img_file)
        shape = img.shape[:2]
        
        gt_file = f"{gt_root}/{name}.xml".replace('R.xml', '_GT.json')
        points = json_points(gt_file)
        density = generate_density_map(shape, points)
    
        fig, ax = plt.subplots()
        ax.imshow(img)
        ax.imshow(density, cmap="jet", alpha=0.7)
        ax.axis('off')

        text_params['s'] = f"GT: {len(points)}"
        text_params['transform'] = ax.transAxes
        ax.text(x=0.5, y=0.02, **text_params) # 距离底部2%的位置（可调节） # 水平居中位置（50%）

        # plt.show()
        plt.savefig(f"{savepath}/{name}_density_{len(points)}.jpg", bbox_inches='tight', pad_inches=0, dpi=300)
        plt.close()
    
        for method, path in results.items():
            if method == 'P2P':
                pre_points = torch.load(os.path.join(path, f'{name}.pt')).detach().cpu().numpy().tolist()
                pre_count = "{:.2f}".format(len(pre_points))
                pre_density = generate_density_map(shape, pre_points, sigma=8)
            elif method == 'FIDTM':
                _, pre_density = get_result(os.path.join(path, f'{name}.pt'), shape)
                pre_count = LMDS_counting(torch.load(os.path.join(path, f'{name}.pt')))
                pre_count = "{:.2f}".format(pre_count)
            else:
                pre_count, pre_density = get_result(os.path.join(path, f'{name}.pt'), shape)
    
            if method == 'HMoDE':
                pre_count = "{:.2f}".format(float(pre_count) / 100.)

            fig, ax = plt.subplots()
            ax.imshow(img)
            ax.imshow(pre_density, cmap="jet", alpha=0.7)
            ax.axis('off')
            
            text_params['s'] = f"ET: {pre_count}"
            text_params['transform'] = ax.transAxes
            ax.text(x=0.5, y=0.02, **text_params) # 距离底部2%的位置（可调节） # 水平居中位置（50%）
            
            plt.savefig(f"{savepath}/{name}_{method}_{pre_count}.jpg", bbox_inches='tight', pad_inches=0, dpi=300)
            plt.close()




cmp_rgb_names = [2800, 2158, 2166]
save_rgbtcc_rgb(cmp_rgb_names)



cmp_rgbt_names = [1565, 1686, 1705]
save_rgbtcc_rgbt(cmp_rgbt_names)

